#!/bin/bash -e

# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Enable recursive globbing
shopt -s globstar

if [ -z "$1" ]; then
    echo "Usage: $0 <directory_containing_wheels>"
    exit 1
fi

# shellcheck disable=2206
WHEEL_LIST=($1/**/cvcuda*.whl)
PACKAGES=("cvcuda")

# Check for GPU and driver
if command -v nvidia-smi &> /dev/null; then
    if ! nvidia-smi > /dev/null 2>&1; then
        echo "Warning: No GPU detected or driver not working with nvidia-smi. Skipping wheel testing..."
        exit 0
    fi
    echo "GPU detected and driver is working (via nvidia-smi)."
elif command -v tegrastats &> /dev/null; then
    if ! tegrastats > /dev/null 2>&1; then
        echo "Warning: tegrastats could not verify GPU. Skipping wheel testing..."
        exit 0
    fi
    echo "GPU detected and driver is working (via tegrastats)."
else
    echo "Warning: Neither nvidia-smi nor tegrastats found. Skipping wheel testing..."
    exit 0
fi

# Check if there are any wheels to test
if [ ${#WHEEL_LIST[@]} -eq 0 ]; then
    echo "No wheels found in the specified directory."
    exit 0
fi

# Track if any wheels were successfully tested
wheels_tested=0

# Extract compatible Python versions from the wheel filenames
get_compatible_pythons() {
    local wheel_file="$1"
    local compatible_versions=()
    python_tag=$(basename "${wheel_file}" | cut -d'-' -f3)
    IFS='.' read -ra tags <<< "${python_tag}"

    for tag in "${tags[@]}"; do
        if [[ "${tag}" =~ cp(3[0-9]{1,2}) ]]; then
            # CPython-specific tag (e.g., cp39, cp310)
            py_ver="${BASH_REMATCH[1]}"
            compatible_versions+=("${py_ver}")
        elif [[ "${tag}" == "py3" ]]; then
            # Generic py3 tag - compatible with all Python 3.x versions
            # Find all available Python 3.x versions on the system
            for py_cmd in python3.{9..13}; do
                if command -v "${py_cmd}" &> /dev/null; then
                    # Extract version number (e.g., "39" from "python3.9")
                    py_ver="${py_cmd##*.}"
                    compatible_versions+=("3${py_ver}")
                fi
            done
            # Break after handling py3 to avoid duplicates
            break
        fi
    done

    echo "${compatible_versions[@]}"
}

# Test each wheel
for whl in "${WHEEL_LIST[@]}"; do
    echo "Testing wheel: $(basename "${whl}")"

    # Determine compatible Python versions
    # shellcheck disable=2207
    compatible_pythons=($(get_compatible_pythons "${whl}"))
    if [ "${#compatible_pythons[@]}" -eq 0 ]; then
        echo "Error: No compatible Python versions found for $(basename "${whl}"). Skipping."
        continue
    fi

    for py_ver in "${compatible_pythons[@]}"; do
        py_exec="python3.${py_ver:1}"
        if ! command -v "${py_exec}" &> /dev/null; then
            echo "Skipping Python ${py_ver}: ${py_exec} not found."
            continue
        fi

        echo "Testing with ${py_exec}..."

        # Create a temporary virtual environment
        test_env_dir=$(mktemp -d)
        ${py_exec} -m venv --without-pip "${test_env_dir}/test_env"
        source "${test_env_dir}/test_env/bin/activate"

        # Manually install pip using ensurepip (should work in venv even on Debian/Ubuntu)
        echo "Manually installing pip..."
        if ! ${py_exec} -m ensurepip --upgrade; then
            echo "Error: Failed to install pip using ensurepip in venv."
            echo "This is unexpected - ensurepip should work in virtual environments."
            deactivate
            rm -rf "${test_env_dir}"
            exit 1
        fi

        # Ensure pip is up to date
        echo "Upgrading pip..."
        ${py_exec} -m pip install --upgrade pip

        # Install and test the wheel
        echo "Installing ${whl}..."
        ${py_exec} -m pip install "${whl}"

        for package in "${PACKAGES[@]}"; do
            echo "Testing import for package ${package}..."
            if ! ${py_exec} <<'PYTEST'
import cvcuda
print('SUCCESS: cvcuda imported')
s = cvcuda.Stream()
print('SUCCESS: Stream created')
t = cvcuda.Tensor((10, 20), cvcuda.Type.U8)
print('SUCCESS: Tensor created')
print('SUCCESS: Version:', cvcuda.__version__)
s.sync()
print('SUCCESS: Stream synced')
PYTEST
            then
                echo "  Error: Failed to import ${package} with ${py_exec}."
                echo "  Please check the error output above for details."
                deactivate
                rm -rf "${test_env_dir}"
                exit 1
            fi
            echo "  ${package} imported successfully with ${py_exec}."
        done


        deactivate
        rm -rf "${test_env_dir}"
        echo "Testing with ${py_exec} completed successfully."

        # Increment counter for successfully tested wheels
        wheels_tested=$((wheels_tested + 1))
    done
done

# Check if any wheels were actually tested
if [ ${wheels_tested} -eq 0 ]; then
    echo "Error: No wheels were successfully tested. All wheels were skipped."
    echo "This could be due to:"
    echo "  - No compatible Python versions installed"
    echo "  - All wheels have invalid Python tags"
    echo "  - Other configuration issues"
    exit 1
fi

echo "All wheels tested successfully. Total tests run: ${wheels_tested}"
